#include <torch/torch.h>
#include <iostream>
#include <vector>
#include <random>
#include <algorithm>  // 用于排序

// 定义前馈神经网络模型
struct FeedForwardNN : torch::nn::Module {
    torch::nn::Linear fc1{nullptr}, fc2{nullptr}, fc3{nullptr};

    FeedForwardNN(int input_size, int hidden_size, int output_size) {
        fc1 = register_module("fc1", torch::nn::Linear(input_size, hidden_size));
        fc2 = register_module("fc2", torch::nn::Linear(hidden_size, hidden_size));
        fc3 = register_module("fc3", torch::nn::Linear(hidden_size, output_size));
    }

    torch::Tensor forward(torch::Tensor x) {
        x = torch::relu(fc1(x));
        x = torch::relu(fc2(x));
        x = fc3(x);
        return x;
    }
};

// 生成数据集（正态分布的 key）
void generate_data(std::vector<float>& keys, std::vector<int>& positions, size_t num_samples) {
    std::normal_distribution<float> key_dist(0.0, 1e10);  // 均值为0，标准差为1e10的正态分布
    std::mt19937 key_gen(time(nullptr));  // 随机数生成器

    keys.resize(num_samples);
    positions.resize(num_samples);

    for (size_t i = 0; i < num_samples; ++i) {
        keys[i] = key_dist(key_gen);  // 生成正态分布的key
        positions[i] = i;  // 使用顺序位置作为目标值
    }
}

// 划分为N个区间
std::vector<std::pair<float, float>> divide_into_intervals(
        const std::vector<std::pair<float, float>>& key_predictions, int N) {

    int num_predictions = key_predictions.size();
    int interval_size = num_predictions / N;

    std::vector<std::pair<float, float>> intervals;

    for (int i = 0; i < N; ++i) {
        float start = key_predictions[i * interval_size].second;  // 取出排序后的预测值
        float end = (i == N - 1) ? key_predictions.back().second : key_predictions[(i + 1) * interval_size - 1].second;
        intervals.push_back({start, end});
    }

    return intervals;
}

int main() {
    size_t num_samples = 1e8;  // 测试数据集大小 1e8
    size_t batch_size = 10000; // 每批次大小
    int num_epochs = 100;      // 训练轮次
    int N = 10;  // 划分为10个区间

    // 生成数据
    std::vector<float> keys;
    std::vector<int> positions;
    generate_data(keys, positions, num_samples);

    // 创建模型和优化器
    auto model = std::make_shared<FeedForwardNN>(1, 64, 1);  // 输入维度1，隐藏层64，输出维度1
    torch::optim::Adam optimizer(model->parameters(), torch::optim::AdamOptions(0.001));

    // 训练模型 - 使用 Mini-Batch
    size_t num_batches = num_samples / batch_size;
    for (int epoch = 0; epoch < num_epochs; ++epoch) {
        model->train();
        float epoch_loss = 0.0;

        for (size_t batch_idx = 0; batch_idx < num_batches; ++batch_idx) {
            // 获取当前批次的数据
            auto start_idx = batch_idx * batch_size;
            auto end_idx = std::min(start_idx + batch_size, num_samples);

            std::vector<float> batch_keys(keys.begin() + start_idx, keys.begin() + end_idx);
            std::vector<int> batch_positions(positions.begin() + start_idx, positions.begin() + end_idx);

            // 转换为张量
            torch::Tensor input_tensor = torch::tensor(batch_keys, torch::kFloat32).unsqueeze(1);
            torch::Tensor target_tensor = torch::tensor(batch_positions, torch::kFloat32).unsqueeze(1);

            // 前向传播
            optimizer.zero_grad();
            auto output = model->forward(input_tensor);
            auto loss = torch::mse_loss(output, target_tensor);

            // 反向传播并优化
            loss.backward();
            optimizer.step();

            epoch_loss += loss.item<float>();
        }

        std::cout << "Epoch [" << epoch << "/" << num_epochs << "], Loss: " << epoch_loss / num_batches << std::endl;
    }

    // 预测所有数据点，保留 key 和预测值 (批量预测)
    std::vector<std::pair<float, float>> key_predictions;
    for (size_t batch_idx = 0; batch_idx < num_batches; ++batch_idx) {
        auto start_idx = batch_idx * batch_size;
        auto end_idx = std::min(start_idx + batch_size, num_samples);

        std::vector<float> batch_keys(keys.begin() + start_idx, keys.begin() + end_idx);

        // 转换为张量并进行预测
        torch::Tensor input_tensor = torch::tensor(batch_keys, torch::kFloat32).unsqueeze(1);
        auto batch_outputs = model->forward(input_tensor);

        for (size_t i = 0; i < batch_outputs.size(0); ++i) {
            key_predictions.push_back({batch_keys[i], batch_outputs[i].item<float>()});
        }
    }

    // 对 key 和预测值按预测值进行排序
    std::sort(key_predictions.begin(), key_predictions.end(),
              [](const std::pair<float, float>& a, const std::pair<float, float>& b) {
                  return a.second < b.second;
              });

    // 将数据划分为 N 个区间
    std::vector<std::pair<float, float>> intervals = divide_into_intervals(key_predictions, N);

    // 输出划分的区间
    std::cout << "\nDivided Intervals:\n";
    for (int i = 0; i < intervals.size(); ++i) {
        std::cout << "Interval " << i + 1 << ": [" << intervals[i].first << ", " << intervals[i].second << "]\n";
    }

    return 0;
}
